#!/usr/bin/env python
# -*- encoding: UTF-8 -*-
from __future__ import print_function
from common import get_line2
import numpy as np
from numpy.linalg import inv
from robot import transl,r2t,t2r,ctraj

__doc__='''
紧协调: 初始位姿我摆出来,则可通过相对位姿确定个机械臂的角度,

每当上双目检测出直线,计算其与棋盘格平面的角点和角度,

计算保持直线和棋盘格相对位姿不变的坐标(棋盘格应该的位姿)

若其位姿与当前的棋盘格位姿的距离大于一个值 e

计算当前的主机器人左臂末端位姿,和目标位姿,进行差值,

对每个差值细分,计算其他协同机械臂的差值,

对每个结果求运动学逆解,可以绕目标位姿的z轴转动,

将计算结果即个关节角度用 api 发送给给机器人.
'''

from sql import Nao
import cv2
import os
import csv
from common import datetime2str
from datetime import datetime

red = Nao.get_by_name("red")
blue = Nao.get_by_name("blue")

class Loger(object):
    def __init__(self):
        self._csv_name = "jinxietiao.csv"
        create_header = not os.path.exists(self._csv_name)
        self._csv = open(self._csv_name,'a')
        self.fieldnames = ['time', 'name','data']
        self.writer = csv.DictWriter(self._csv, fieldnames=self.fieldnames)
        if create_header:
            self.writer.writeheader()

    def log(self,**kwargs):
        log_dict = {}
        if 'time' not in kwargs:
            log_dict['time'] = str(datetime.now())
        for k,v in kwargs.iteritems():
            if k in self.fieldnames:
                log_dict[k]=v
        self.writer.writerow(log_dict)

    def _get_reader(self):
        self._csv_read =  open(self._csv_name,'r')
        self.reader = csv.DictReader(self._csv_read)

    def get(self,fieldname):
        if not hasattr(self,'reader'):
            self._get_reader()
        out_list = []
        for row in self.reader:
            out_list.append(row[fieldname])
        return out_list

class Cooperation(object):
    def __init__(self,master,slave):
        assert isinstance(master,Nao)
        assert isinstance(slave,Nao)
        self.master = master
        self.slave = slave
        self.w = 9 # 棋盘格宽
        self.h = 4 # 棋盘格高
        self.unit = 15 # 棋盘格单位长度 mm
        self.dt = 0.01 # 仿真精度
        self.Tmct_mcb = self.master.tr_top_bottom()
        self.Tsct_scb = self.slave.tr_top_bottom()

    def init_pose(self):
        '''
        读取各关节角度,求解四个机械手末端位姿与标定板坐标系的相对位姿,
        计算空间直线与标定板平面(xy)的交点在标定板坐标系下的表示
        因为该函数依赖于 tf 故要先运行 get_image,get_tf
        :return:
        '''
        assert self.Tmcb_obj is not None
        assert self.Tscb_obj is not None
        self.initTobj_ml = np.linalg.inv(self.Tmo_obj)*self.Tmo_ml
        self.initTobj_mr = np.linalg.inv(self.Tmo_obj)*self.Tmo_mr
        self.initTobj_sl = np.linalg.inv(self.Tso_obj)*self.Tso_sl
        self.initTobj_sr = np.linalg.inv(self.Tso_obj)*self.Tso_sr
        self.initPoint,self.initVec = self.get_line_point_and_vec_in_obj()
        assert self.initPoint is not None
        assert self.initVec is not None
        self.check_point(self.initPoint)
        x,y = self.initPoint[:2]
        # {obj} 为 opencv 棋盘格坐标系,以左上角为原点
        # {line_obj} 为棋盘格上以直线与 Oxy 平面的交点为原点的坐标系
        self.initTobj_line_obj = transl(x,y,0)
        self.init_obj_mz = self._get_obj_mz(self.Tmo_obj,self.initPoint)
        self.weight_x = np.linalg.norm([self.initTobj_ml[0,3] - x,
                                        self.initTobj_mr[0,3] - x,
                                        self.initTobj_sl[0,3] - x,
                                        self.initTobj_sr[0,3] - x
                                        ])
        self.weight_y = np.linalg.norm([self.initTobj_ml[1,3] - y,
                                        self.initTobj_mr[1,3] - y,
                                        self.initTobj_sl[1,3] - y,
                                        self.initTobj_sr[1,3] - y,
                                        ])

    def get_image(self):
        self._master_top_image = self.master.top_cam.get_image()
        self._master_bottom_image = self.master.bottom_cam.get_image()
        self._slave_top_image = self.slave.top_cam.get_image()
        self._slave_bottom_image = self.slave.bottom_cam.get_image()

    def get_tf(self):
        '''
        因为该函数依赖于图像故要先运行 get_image
        :return:
        '''
        # Tmo_ml : 在 master 机器人坐标系下 LArm 的末端位姿
        self.Tmo_ml = self.master._getTf("LArm")
        self.Tmo_mr = self.master._getTf("RArm")
        self.Tso_sl = self.slave._getTf("LArm")
        self.Tso_sr = self.slave._getTf("RArm")
        # Tmo_ct: 在 master 机器人坐标系下 camera top 的位姿
        self.Tmo_ct = self.master._getTf("TopCamera")
        self.Tmo_cb = self.master._getTf("BottomCamera")
        self.Tso_ct = self.slave._getTf("TopCamera")
        self.Tso_cb = self.slave._getTf("BottomCamera")
        # Tmcb_obj: 在 master 机器人下相机下标定板的位姿,由相机标定函数求得
        self.Tmcb_obj = self.master.bottom_cam.findTr(self,self._master_bottom_image,
                                                      w=self.w,h=self.h,unit=self.unit)
        self.Tscb_obj = self.slave.bottom_cam.findTr(self,self._slave_bottom_image,
                                                      w=self.w,h=self.h,unit=self.unit)
        if self.Tmcb_obj is None or self.Tscb_obj is None:
            self.Tso_mo,self.Tmct_obj,self.Tscb_obj,self.Tmo_obj,self.Tso_obj = None,None,None,None,None
        else:
            self.Tmct_obj = self.Tmct_mcb*self.Tmcb_obj
            self.Tsct_obj = self.Tsct_scb*self.Tscb_obj
            # Tmo_obj: 在 master 机器人坐标系下 目标(标定板)的位姿
            self.Tmo_obj = self.Tmo_cb*self.Tmcb_obj
            self.Tso_obj = self.Tso_cb*self.Tscb_obj
            self.Tso_mo = self.Tso_obj*inv(self.Tmo_obj)

        self.get_line_from_top_image()

    def get_line_from_top_image(self):
        '''
        在 obj 坐标系下直线的 Plucker 表示,
        :return: L 两点式表示, Lstart 两面式表示
        '''
        if self.Tmcb_obj is None or self.Tscb_obj is None: return None
        L,Lstart = get_line2(self._master_top_image,
                             self.master.top_cam,
                             self.Tmct_obj,
                             self._slave_top_image,
                             self.slave.top_cam,
                             self.Tsct_obj)
        self.L = L
        self.Lstart = Lstart

    def _unit(self,w):
        w[:3,0] = w[:3,0]/np.linalg.norm(w[:3,0])
        return w

    def _get_obj_mz(self,Tmo_obj,point):
        '''
        获取指定obj的Oxy平面上一点在 master 机器人坐标系下的 z 值
        :param Tmo_obj:  {}_{obj}^{mo} T 指定obj在 master 机器人坐标系下的位姿
        :param point:  在指定 obj 下 Oxy 平面上一点的坐标 (x,y)
        :return: mz
        '''
        x,y = np.array(point).flatten()[:2]
        Tmo_lineobj = Tmo_obj*transl(x,y,0)
        return Tmo_lineobj[2,3]

    def _set_obj_mz(self,Tmo_obj,point,mz,vec):
        '''
        设定 指定obj的Oxy平面上一点在 master 机器人坐标系下的 z 值
        :param Tmo_obj: 指定obj在 master 机器人坐标系下的位姿
        :param point: 在指定 obj 下 Oxy 平面上一点的坐标 (x,y)
        :param mz: 设定的 z 值
        :param vec: 在 master 坐标系下用来平移的方向向量
        :return: new_Tmo_obj 新的坐标系
        '''
        old_obj_mz = Tmo_obj[2,3]
        old_mz = self._get_obj_mz(Tmo_obj,point)
        delta = old_obj_mz - old_mz
        obj_mz = delta + mz
        obj_pos = Tmo_obj[:3,3]
        k = (obj_mz - old_obj_mz)/vec[2,0]
        new_obj_pos = obj_pos + k*vec[:3,0]
        new_Tmo_ojb = Tmo_obj.copy()
        new_Tmo_ojb[:3,3] = new_obj_pos
        return new_Tmo_ojb

    def get_Tmo_newobj(self):
        '''
        根据初始条件,保持 z 轴与检测出的直线平行,即使检测出来的直线与标定板保持垂直,
        并保持新的坐标系的指定点在 mo 坐标系下的z轴值不变
        :return:
        '''
        point,vec = self.get_line_point_and_vec_in_obj()
        if point is None or vec is None: return None
        if self.Tmo_obj is None: return None
        R_newobj_obj = self._get_R_nowobj_obj(vec)
        x,y = np.array(self.initPoint).flatten()[:2]
        dx,dy = np.array(point-self.initPoint).flatten()[:2]
        new_Tmo_new_obj = self.Tmo_obj*transl(x,y,0)*transl(dx,dy,0)*r2t(inv(R_newobj_obj))*transl(-x,-y,0)
        new_Tmo_new_obj = self._set_obj_mz(new_Tmo_new_obj,self.initPoint,self.init_obj_mz,self.Tmo_obj*vec)
        return new_Tmo_new_obj

    def get_line_point_and_vec_in_obj(self):
        '''
        计算在 obj 坐标系下,空间直线与 Oxy 平面的交点,即目标点
        :return:
        '''
        L,Lstar = self.L,self.Lstart
        if L is None or Lstar is None: return None,None
        # 求直线方向向量
        vec = L[:3,3]
        vec = vec/np.linalg.norm(vec)
        if vec[2] < 0:
            vec = - vec # 保持直线的方向向量沿z轴正方向
        vec = np.vstack((vec,1)) # 齐次化
        vec = np.mat(vec).reshape((4,1))
        # 求与 Oxy 平面的交点
        S,V,D = np.linalg.svd(Lstar)
        D = D.transpose()
        p1 = D[:,2] # 直线上的一个点
        p1 = p1/p1[-1,0] # 齐次化
        p2 = D[:,3] # 直线上的另一个点
        p2 = p2/p2[-1,0] # 齐次化
        point = p1 - vec*p1[2]/vec[2] # 与 Oxy 平面的交点
        point = point/point[-1,0] # 齐次化
        return point,vec

    def _get_R_nowobj_obj(self,vec):
        new_z = vec[:3,0]
        a,b,c = np.array(vec).flatten()[:3]
        # 原 x 轴单位向量在与 new_z 垂直平面的投影向量
        proj_x = np.mat([1-a**2, - a*b, - a*c]).reshape((3,1))
        proj_x = proj_x/np.linalg.norm(proj_x) # 单位化
        # 原 y 轴单位向量在与 new_z 垂直平面的投影向量
        proj_y = np.mat([-a*b,1-b**2,-b*c]).reshape((3,1))
        proj_y = proj_y/np.linalg.norm(proj_y) # 单位化
        # 两个轴的夹角
        dot = (proj_x.T*proj_y)[0,0]
        theta0 = np.arccos(dot)
        # x轴投影向量作为新的 x 轴向量则与原 x 轴的夹角最小
        # 综合考虑,将新的x轴向量绕 新的z轴转动一个角度
        theta1 = (theta0 - np.pi/2)/(1+self.weight_x/self.weight_y)
        new_x = np.mat(cv2.Rodrigues(vec*theta1)[0])*proj_x
        new_y = np.cross(new_z,new_x,axis=0)
        R_newobj_obj = np.mat(np.hstack((new_x,new_y,new_z))).T
        return R_newobj_obj

    def check_point(self,point):
        '''
        确保空间直线与标定板平面的交点在合理范围内
        :param point:
        :return:
        '''
        x,y = point[:2]
        assert  0 <= x <= self.w*self.unit
        assert  0 <= y <= self.h*self.unit

    def run(self):
        self.get_image()
        self.get_tf()
        self.init_pose()
        while True:
            self.get_image()
            self.get_tf()
            Tmo_newobj = self.get_Tmo_newobj()
            if self.delta_Tf(Tmo_newobj,self.Tmo_obj) > 0.1:
                self.move_to(Tmo_newobj)

    def delta_Tf(self,Tf1,Tf2):
        delta = Tf1-Tf2
        delta[:3,3] *= 0.001
        return np.linalg.norm(delta)

    def move_to(self,Tmo_newobj,v=0.1):
        Tmo_obj_list,time = self.get_move_Tmo_obj_list(self.Tmo_obj,Tmo_newobj,v=v)
        Tmo_ml_list = self.get_Tmo_ml_list_from_Tmo_obj_list(Tmo_obj_list)
        Tmo_mr_list = self.get_Tmo_mr_list_from_Tmo_obj_list(Tmo_obj_list)
        Tso_obj_list = self.get_Tso_obj_list_from_Tmo_obj_list(Tmo_obj_list)
        Tso_sl_list = self.get_move_Tso_sl_list_from_Tso_obj_list(Tso_obj_list)
        Tso_sr_list = self.get_move_Tso_sr_list_from_Tso_obj_list(Tso_obj_list)
        ml_angles = self.get_ml_angles_from_Tmo_ml_list(Tmo_ml_list)
        mr_angles = self.get_mr_angles_from_Tmo_mr_list(Tmo_mr_list)
        sl_angles = self.get_sl_angles_from_Tso_sl_list(Tso_sl_list)
        sr_angles = self.get_sr_angles_from_Tso_sr_list(Tso_sr_list)
        self.move_angles(time,ml_angles,mr_angles,sl_angles,sr_angles)

    def get_move_Tmo_obj_list(self,Tmo_obj,Tmo_newobj,v=0.1):
        return self.ctraj(Tmo_obj,Tmo_newobj,v=v)

    def ctraj(self,T0,T1,v=0.1):
        d = self.delta_Tf(T0,T1)
        #time = d/v
        time = 1.0 # 每一次检测出来就运动 1s
        return ctraj(T0,T1,time/self.dt)

    def _mul_left(self,A,B_list):
        out_list = []
        for B in B_list:
            out_list.append(A*B)
        return out_list

    def _mul_right(self,A_list,B):
        out_list = []
        for A in A_list:
            out_list.append(A*B)
        return out_list

    def get_Tmo_ml_list_from_Tmo_obj_list(self,Tmo_obj_list):
        return self._mul_right(Tmo_obj_list,self.initTobj_ml)

    def get_Tmo_mr_list_from_Tmo_obj_list(self,Tmo_obj_list):
        return self._mul_right(Tmo_obj_list,self.initTobj_mr)

    def get_Tso_obj_list_from_Tmo_obj_list(self,Tmo_obj_list):
        return self._mul_left(self.Tso_mo,Tmo_obj_list)

    def get_move_Tso_sl_list_from_Tso_obj_list(self,Tso_obj_list):
        return self._mul_right(Tso_obj_list,self.initTobj_sl)

    def get_move_Tso_sr_list_from_Tso_obj_list(self,Tso_obj_list):
        return self._mul_right(Tso_obj_list,self.initTobj_sr)

    def _get_angles_from_Tr_list(self,Tr_list,q0,robot):
        angles = []
        _q_0 = np.array(q0)
        new_q = _q_0.copy()
        for Tr in Tr_list:
            new_q[:-1] = robot.ikine2(Tr,_q_0[:-1])
            new_q[-1] = _q_0[-1]
            angles.append(new_q.copy())
            _q_0 = new_q
        return np.array(angles)

    def get_ml_angles_from_Tmo_ml_list(self,Tmo_ml_list):
        Tr_list = Tmo_ml_list
        q0 = self.master._getAngles("LArm")
        robot = self.master.kin.LeftHand
        return self._get_angles_from_Tr_list(Tr_list,q0,robot)

    def get_mr_angles_from_Tmo_mr_list(self,Tmo_mr_list):
        Tr_list = Tmo_mr_list
        q0 = self.master._getAngles("RArm")
        robot = self.master.kin.RightHand
        return self._get_angles_from_Tr_list(Tr_list,q0,robot)

    def get_sl_angles_from_Tso_sl_list(self,Tso_sl_list):
        Tr_list = Tso_sl_list
        q0 = self.slave._getAngles("LArm")
        robot = self.slave.kin.LeftHand
        return self._get_angles_from_Tr_list(Tr_list,q0,robot)

    def get_sr_angles_from_Tso_sr_list(self,Tso_sr_list):
        Tr_list = Tso_sr_list
        q0 = self.slave._getAngles("RArm")
        robot = self.slave.kin.RightHand
        return self._get_angles_from_Tr_list(Tr_list,q0,robot)

    def move_angles(self,time,ml_angles,mr_angles,sl_angles,sr_angles):
        names = [ 'LShoulderPitch',
                  'LShoulderRoll',
                  'LElbowYaw',
                  'LElbowRoll',
                  'LWristYaw',
                  'LHand',
                  'RShoulderPitch',
                  'RShoulderRoll',
                  'RElbowYaw',
                  'RElbowRoll',
                  'RWristYaw',
                  'RHand']
        length = len(ml_angles)
        timesList = np.arange(0.0,time,time/length)
        timesList = np.array([timesList for n in xrange(12)])
        timesList = timesList.tolist()
        master_angles = np.hstack((ml_angles,mr_angles))
        master_angles_Lists = master_angles.transpose().tolist()
        slave_angles = np.hstack((sl_angles,sr_angles))
        slave_angles_Lists = slave_angles.transpose().tolist()
        self.master.angleInterpolation(names,master_angles_Lists,timesList)
        self.slave.angleInterpolation(names,slave_angles_Lists,timesList)

